from src.dataset import PolyvoreTripletDataset, CompatibilityBenchmarkDataset, FITBBenchmarkDataset
from src.const import base_path
import matplotlib.pyplot as plt
import pandas as pd
import torch
import torch.utils.data
from torch import nn
import numpy as np
from torch.nn import functional as F
from src.const import base_path
from src import const
from src.utils import parse_args_and_merge_const, load_json, build_vocab
from tensorboardX import SummaryWriter
import os


if __name__ == '__main__':
    args = parse_args_and_merge_const()
    assert (args.model != "")
    if os.path.exists('models') is False:
        os.makedirs('models')
    folder_name = os.path.join('models', const.MODEL_NAME)
    if not(os.path.exists(folder_name)):
        os.makedirs(folder_name)
    
    # 强制使用大数据集来evaluate
    const.VAL_FASHION_COMP_FILE = "fashion_compatibility_prediction.txt"
    const.VAL_FITB_FILE = "fill_in_blank_test.json"
    test_set = load_json(os.path.join(const.base_path, 'test_no_dup.json'))
    comp_dataset = CompatibilityBenchmarkDataset(const.DATASET_PROC_METHOD_VAL, test_set)
    fitb_dataset = FITBBenchmarkDataset(const.DATASET_PROC_METHOD_VAL, test_set)

    net = const.USE_NET(pretrained_embeddings=None)
    net = net.to(const.device)
    net.load_state_dict(torch.load(args.model))
    
    print('Now Evaluate..')
    with torch.no_grad():
        net.eval()
        fitb_benchmark = fitb_dataset.calculate(net)
        print('FITB Benchmark: {}'.format(fitb_benchmark))
        comp_benchmark = comp_dataset.calculate(net)
        print('Compatibility Benchmark: {}'.format(comp_benchmark))
